# vim: syntax=cuda
# Kernels for bootstrapping rows and columns

"""

Copyright 2009 Michael Seiler
Rutgers University
miseiler@gmail.com

This file is part of ConsensusCluster.

ConsensusCluster is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

ConsensusCluster is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with ConsensusCluster.  If not, see <http://www.gnu.org/licenses/>.


"""

from __future__ import division

import cu_twister      as curand
import pycuda.driver   as cuda
import pycuda.gpuarray as gpuarray
import numpy           as N

BS_ROW_ROWS = 16
BS_ROW_COLS = 32

from kernel_helpers import Krnl

# Global vars
MAX_RAND_SEED = 32768

bootstrap_rows_kernel = """

    // Bootstrap rows of a matrix

    // oldrows is the number of rows before padding
    // New dim requests are newrows
    // out should be <new row dim> x cols
    // Both in and out need to have dims multiple to BLOCK_ROWS/BLOCK_COLS
    // Grid this in a fully-tiled fashion according to the dims of out
    
    #define BLOCK_ROWS %(brows)s
    #define BLOCK_COLS %(bcols)s

    #include "mt_rand.cu.h"

    __global__ void bootstrap (float *out, float *in, int *samples, int cols, int newrows, int oldrows) {
    	
        int tid       = __mul24(threadIdx.y, BLOCK_COLS) + threadIdx.x;
        int blk_start = (__mul24(BLOCK_ROWS, blockIdx.y) + threadIdx.y) * cols;
        int xoffset   = __mul24(blockIdx.x, BLOCK_COLS) + threadIdx.x;
        int ridx      = __mul24(blockIdx.y, BLOCK_ROWS) + tid;
        
        // Current twister state

        __shared__ MersenneTwisterState rstate[BLOCK_ROWS];

        // Indices of next sample

        __shared__ int idx[BLOCK_ROWS];
        
        int iters, currow, out_row;

        if (tid < BLOCK_ROWS)
            MersenneTwisterInitialise(rstate[tid], ridx);

        // Unfortunately, this is quite a register-inefficient way to determine number of iterations
        iters = newrows / BLOCK_ROWS;

        if (iters > gridDim.y) {
            iters = (iters / gridDim.y) + (blockIdx.y < iters %% gridDim.y);
        } else {
            iters = (blockIdx.y < iters);
        }

        while (iters--) {

            if (tid < BLOCK_ROWS)
                idx[tid] = (int) (mt_rand(rstate[tid], ridx) * oldrows) * cols;

            __syncthreads();
    
            currow = idx[threadIdx.y];
    
            __syncthreads();

            out_row = blk_start + iters * gridDim.y * cols;

            out[ out_row + xoffset ] = in[ currow + xoffset ];
    
            __syncthreads();
            
            if (blockIdx.x == 0 && tid < BLOCK_ROWS)
                samples[ __mul24(BLOCK_ROWS, blockIdx.y) + iters * gridDim.y + tid ] = idx[tid] / cols;

        }
    }
"""


def get_bootstrap_func():

    int32 = N.int32
    seed = curand.seed

    max_init_blocks = MAX_RAND_SEED // BS_ROW_ROWS

    _bootstrap              = Krnl(bootstrap_rows_kernel, (BS_ROW_COLS, BS_ROW_ROWS), frmt = {'brows': BS_ROW_ROWS, 'bcols': BS_ROW_COLS}, include_dirs=[curand.get_include_dir()])
    _bootstrap_f, block_bs  = _bootstrap.func('bootstrap', 'PPPiii')

    def bootstrap(out, inp, samples, oldrows, newrows):

        c = inp.shape[1]
        grid = (c // BS_ROW_COLS, min(newrows // BS_ROW_ROWS, max_init_blocks))

        seed(cuda, _bootstrap.mod)
        _bootstrap_f(grid, block_bs, out.gpudata, inp.gpudata, samples.gpudata, c, int32(newrows), int32(oldrows))

    return bootstrap
